#!/usr/bin/env python
import roslib; roslib.load_manifest('block_build_manager')
roslib.load_manifest('rospy')
roslib.load_manifest('actionlib')
roslib.load_manifest('pr2_controllers_msgs')

import rospy
import actionlib
import traceback

from block_build_msgs.msg import *
from object_manager.msg import *
from object_manager.srv import *
from geometry_msgs.msg import *
from std_msgs.msg import *
from actionlib_msgs.msg import *
from pr2_controllers_msgs.msg import *

from std_srvs import *

""" 
  Author : Jihoon Lee
  Date   : Feb. 2012

  BuildManager.py
"""
COMMAND_NAME = "/blocknlp/command"
PICK_COMMAND_NAME = "/blocknlp/pick"
PLACE_COMMAND_NAME = "/blocknlp/place"
OBJECT_LIST_TOPIC = "/mit/object_list_all"
OBJECT_UPDATE_TOPIC_NAME = "/blocknlp/update_object"
OBJECT_UPDATE_ATHOME_TOPIC_NAME = "/blocknlp/update_object_atHome"
UPDATE_TABLE_TOPIC_NAME = "/blocknlp/update_table"
SLU_TOPIC_NAME = "/blocknlp/nl_output"
SLU_COMMAND_PARSER_STATUS_TOPIC_NAME = "/blocknlp/slu_command_parser_status"

class BuildManager:

  
  def __init__(self):
    self.started = False
    self.inProcessCommand = False
    self.object_inhand = None


    self.command_server = actionlib.SimpleActionServer(COMMAND_NAME,block_build_msgs.msg.CommandAction,self.commandCallback,False)
    self.pick_client = actionlib.SimpleActionClient(PICK_COMMAND_NAME,block_build_msgs.msg.PickCommandAction)
    self.place_client = actionlib.SimpleActionClient(PLACE_COMMAND_NAME,block_build_msgs.msg.PlaceCommandAction)

    self.slu_command_listener = rospy.Subscriber(SLU_TOPIC_NAME,SLUCommand,self.sluCommandListener)

    self.object_update_pub = rospy.Publisher(OBJECT_UPDATE_TOPIC_NAME,Object)
    self.table_update_pub = rospy.Publisher(UPDATE_TABLE_TOPIC_NAME,Empty)
    self.slu_command_parser_status_pub = rospy.Publisher(SLU_COMMAND_PARSER_STATUS_TOPIC_NAME, String)

  def processCommand(self,req):
    if req.command == 0:
      self.pickup_by_id(req.obj_num)
    elif req.command == 1:
      self.putdown(0,[req.x,req.y])
    elif req.command == 2:
      self.initializeRobot()
    return True

  def initializeRobot(self):
    g = PickCommandGoal()
    g.number = -3

    rospy.loginfo("initializeRobot: Send initialize robot message")
    self.pick_client.send_goal(g)
    self.pick_client.wait_for_result()

    if self.pick_client.get_result().result:
      rospy.loginfo("initializeRobot: Complete")
    else:
      rospy.loginfo("initializeRobot: Failed")

  def slu_command_parser(self,msg):
    putdown = self.putdown
    pickup = self.pickup_by_id
    dropblock = self.drop_block
    stackon = self.stack_on_id
    returnBlockToHomePosition = self.returnBlockToHomePosition
    resetToolbox = self.resetToolbox
    initializeToolbox = self.initializeToolbox
    initializeRobot = self.initializeRobot
    noop = lambda: True
    try:
      eval(msg.command)
    except:
      traceback.print_exc()
      raise
    finally:
      if self.object_inhand == None and msg.do_reset:
        self.initializeRobot()
      self.slu_command_parser_status_pub.publish(msg.id)
    return msg

  def pickup_by_id(self, id):
    # Create the goal object
    g = PickCommandGoal()
    g.number = -2

    obj_list = rospy.wait_for_message(OBJECT_LIST_TOPIC,ObjectList,10.0)

    foundObject = False
    for obj in obj_list.object_list:
      if id == obj.id:
        g.object = obj
        # Temporary fix for orientation problem
        g.object.pose.orientation.x = 0;
        g.object.pose.orientation.y = 0;
        g.object.pose.orientation.z = 0;
        g.object.pose.orientation.w = 1;
        foundObject = True
        break
   
    if not foundObject:
        return False

    self.inProcessCommand = True

    rospy.loginfo("pickup_by_id: Send pick goal")
    self.pick_client.send_goal(g)
    self.pick_client.wait_for_result()
    if self.pick_client.get_result().result:
      rospy.loginfo("pickup_by_id: Complete")
      new_obj = Object()
      new_object = g.object
      new_object.header.frame_id = "/r_gripper_tool_frame"
      new_object.inhand = True
      new_object.atHome = False
      new_object.pose = Pose()
      self.object_inhand = new_object
      self.object_update_pub.publish(new_object)
    else:
      rospy.loginfo("pickup_by_id: Fail")
    self.inProcessCommand = False

  def pickup(self,num):
    g = PickCommandGoal()
    g.number = -2

    obj_list = rospy.wait_for_message(OBJECT_LIST_TOPIC,ObjectList,10.0)

    self.inProcessCommand = True

    if num > len(obj_list.object_list):
      return False
    else:
      g.object = obj_list.object_list[num]

    rospy.loginfo("pickup: Send pick goal")
    self.pick_client.send_goal(g)
    self.pick_client.wait_for_result()
    if self.pick_client.get_result().result:
      rospy.loginfo("pickup: Complete")
      new_obj = Object()
      new_object = g.object
      new_object.header.frame_id = "/r_gripper_tool_frame"
      new_object.inhand = True
      new_object.atHome = False
      new_object.pose = Pose()
      self.object_inhand = new_object
      self.object_update_pub.publish(new_object)
    else:
      rospy.loginfo("pickup: Fail")
    self.inProcessCommand = False

  def putdown(self,id,location):
    #if self.object_inhand == None:
    #  rospy.loginfo("No object in hand")
    #  return True
    if self.object_inhand == None:
      rospy.loginfo("putdown: No object in hand. Picking up object by id.")
      self.pickup_by_id(id)

    self.inProcessCommand = True
    placeGoal = PlaceCommandGoal()
    placeGoal.x = location[0]
    placeGoal.y = location[1]
    placeGoal.z = 0.66
    placeGoal.orientation = Quaternion(0,0.707,0,0.707)
    placeGoal.precise_place = True            
    print(str(placeGoal))
    
    rospy.loginfo("putdown: Send place goal")
    self.place_client.send_goal(placeGoal)
    self.place_client.wait_for_result()
    if self.place_client.get_result().result:
      rospy.loginfo("putdown: Complete")
    else:
      rospy.loginfo("putdown: Fail")

    if self.object_inhand != None:
      self.object_inhand.inhand = False
      self.object_update_pub.publish(self.object_inhand)
      self.object_inhand = None

    self.inProcessCommand = False

  def stack_on_id(self,id):
    if self.object_inhand == None:
      rospy.loginfo("stack_on_id: No object in hand")
      return True

    obj_list = rospy.wait_for_message(OBJECT_LIST_TOPIC,ObjectList,10.0)  
    foundObject = False
    for obj in obj_list.object_list:
      if id == obj.id:
        obj_height = 0.044
        placeGoal = PlaceCommandGoal()
        placeGoal.x = obj.pose.position.x
        placeGoal.y = obj.pose.position.y
        placeGoal.z = obj.pose.position.z + obj_height
        placeGoal.orientation = Quaternion(0,0.707,0,0.707)
        placeGoal.precise_place = False           
        print(str(placeGoal))
        foundObject = True
        break
   
    if not foundObject:
        return False

    self.inProcessCommand = True
    
    rospy.loginfo("stack_on_id: Send place goal")
    self.place_client.send_goal(placeGoal)
    self.place_client.wait_for_result()
    if self.place_client.get_result().result:
      rospy.loginfo("stack_on_id: Complete")
    else:
      rospy.loginfo("stack_on_id: Fail")

    self.object_inhand.inhand = False
    self.object_update_pub.publish(self.object_inhand)
    self.object_inhand = None

    self.inProcessCommand = False

  # Returns the currently held block back to its home position
  def returnBlockToHomePosition(self):
    if self.object_inhand == None:
      rospy.loginfo("returnBlockToHomePosition: No object in hand")
      return True
    
    location = [self.object_inhand.homePosition.position.x, self.object_inhand.homePosition.position.y]
    self.object_inhand.atHome = True
    self.putdown(0,location)

  def drop_block(self):
    if self.object_inhand == None:
      rospy.loginfo("returnBlockToHomePosition: No object in hand")
      return True

    self.inProcessCommand = True
    client = actionlib.SimpleActionClient('r_gripper_controller/gripper_action', Pr2GripperCommandAction)
    client.wait_for_server()

    rospy.loginfo("returnBlockToHomePosition: Attempting to Drop Box")
    rospy.loginfo("returnBlockToHomePosition: Send Open Goal")
    client.send_goal(Pr2GripperCommandGoal(
	Pr2GripperCommand(position = 0.08, max_effort = -1)))
    client.wait_for_result()

    if client.get_state() == GoalStatus.SUCCEEDED:
        rospy.loginfo("returnBlockToHomePosition: Complete")
    else:
        rospy.loginfo("returnBlockToHomePosition: Fail")

    self.object_inhand.inhand = False
    self.object_update_pub.publish(self.object_inhand)
    self.object_inhand = None

    self.inProcessCommand = False

    self.initializeRobot()

  def initializeToolbox(self):
    self.inProcessCommand = True
    obj_list = rospy.wait_for_message(OBJECT_LIST_TOPIC,ObjectList,10.0)

    for obj in obj_list.object_list:
      print "obj", obj, obj.label
      if obj.atHome:
        self.pickup_by_id(obj.id)
        self.returnBlockToHomePosition()

    self.inProcessCommand = False

  def resetToolbox(self):
    self.inProcessCommand = True

    processed_ids = set()
    if self.object_inhand != None:
      processed_ids.add(self.object_inhand.id)
      self.returnBlockToHomePosition()
        

    obj_list = rospy.wait_for_message(OBJECT_LIST_TOPIC,ObjectList,10.0)

    for obj in obj_list.object_list:
      if not obj.atHome and not obj.id in processed_ids: 
        self.pickup_by_id(obj.id)
        self.returnBlockToHomePosition()
        processed_ids.add(obj.id)

    self.inProcessCommand = False

  def sluCommandListener(self,msg):
    rospy.loginfo("sluCommandListener: in SLU command listener msg :")
    rospy.loginfo(msg.command)
    if self.started:
      rospy.loginfo("Executing command.")
      self.slu_command_parser(msg)
    else:
      rospy.loginfo("Ignoring command; not started.")
    
  def commandCallback(self,goal):
    if goal.command == 99:
      rospy.loginfo("commandCallback: It works!")
    else:
      self.processCommand(goal)

    r = CommandResult()
    r.result = True
    self.command_server.set_succeeded(r)


  def spin(self):
    self.command_server.start()
    rospy.loginfo('spin: Wait for command server')
    while (not self.command_server.action_server.started and 
           not self.pick_client.action_server.started and
           not self.place_client.action_server.started):
      rospy.loginfo("sleeping until active")
      rospy.sleep(1)

    self.started = True

    rospy.loginfo('spin: initialized')

    while not rospy.is_shutdown():
      if not self.inProcessCommand:
        t = Empty()
        #self.table_update_pub.publish()
      rospy.sleep(2)


if __name__ == '__main__':
  rospy.init_node('block_build_manager')

  bm = BuildManager()
  bm.spin()

